#!/usr/bin/env python
# coding=utf-8
import sys

sys.path.insert(0, '/home/center/fdse_coffee/src/jaka_pro')
sys.path.insert(0, '/home/center/fdse_coffee/src/jaka_pro/lib')
import math
import socket
import traceback
import time
import roslibpy
#from lib.vision_package.tag_detector import Image_Receiver
import numpy as np
import jkrc
# import rospy
import cv2
from lib.vision_package.tag_detector import Apriltag_Detector
from lib.vision_package.hough_detector import Hough_Detector

ABS = 0
INCR = 1
ERROR = 0
from copy import deepcopy

NUM = [0, 0]
DEBUG = False



class Gripper(object):
    def __init__(self, gripper_ip='192.168.1.29', port=8888):
        super(Gripper, self).__init__()
        # self.ip = gripper_ip
        self.tcp_server_addr = (gripper_ip, port)
        self.sock = None
        self.connect()
        

    def connect(self):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        # keep alive: the time of 'reconnect'; time interval; the count of retry
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
        self.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 10)
        self.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 3)
        self.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 5)
        try:
            self.sock.connect(self.tcp_server_addr)
            print("connect to gripper successfully")
        except Exception:
            print(traceback.format_exc())

    @staticmethod
    def init(ip='192.168.1.29', port=8888):
        init_str = 'FFFEFDFC010802010000000000FB'
        check_init_str = 'FFFEFDFC010801000000000000FB'

        tcp_server_addr = (ip, port)
        sock_gripper = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        # sock_gripper.setsockopt(socket.SOL_SOCKET, socket.SO_DEBUG, 1)

        # sock_gripper.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
        try:
            sock_gripper.connect(tcp_server_addr)
            print("initial: connect to gripper successfully")
        except Exception:
            print(traceback.format_exc())
        sock_gripper.sendall(init_str.encode())
        time.sleep(3)
        recvdata_init = sock_gripper.recv(2048).decode()
        sock_gripper.sendall(check_init_str.encode())
        recvdata_init_check = sock_gripper.recv(2048).decode()
        if (recvdata_init_check == "FFFEFDFC0108010000A5000000FB"):
            print("gripper initial successfully")
            sock_gripper.close()
            return 1
        else:
            print("fail to initial gripper")
            sock_gripper.close()
            return 0

    # (20~100)
    def pose(self, target, new_order=False):
        while True:
            try:
                gripper_pos = str(hex(target)[2:4]).upper()
                grasp_str = 'FFFEFDFC0106020100' + gripper_pos + '000000FB'
                time.sleep(2)
                self.sock.sendall(grasp_str.encode())
                recvdata_grasp = self.sock.recv(2048).decode()
                time.sleep(2)
                break
            except socket.error as msg:
                print(msg)
                print('Gripper meets a BIG Error and retry ...')
                self.close()
                if new_order:
                    self.init()
                self.connect()
                print('reconnect')

    def open(self):
        self.pose(100)

    def close(self):
        # self.sock.shutdown(socket.SHUT_RDWR)
        self.sock.close()

fx = 610.33349609375
fy = 608.6472778320312
cx = 323.9746398925781
cy = 237.62242126464844
camera_intrinsics = [fx, fy, cx, cy]
detector_tag = Apriltag_Detector(intrinsics=camera_intrinsics)

class Jaka(object):
    def __init__(self, jaka_ip='192.168.1.133'):
        super(Jaka, self).__init__()
        self.ip = jaka_ip
        self.robot = jkrc.RC(self.ip)
        self.robot.login()
        self.gripper = Gripper()
        # rospy.init_node("my_camera")
        
        k = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
        # self.detector_tag = Apriltag_Detector(intrinsics=self.camera_intrinsics)
        self.detector_hough = Hough_Detector()
        # detector_tag = Apriltag_Detector(intrinsics=camera_intrinsics)
        self.cap = cv2.VideoCapture(0)
        #self.image_recv = Image_Receiver()
        self.holder_pos = -1
        self.work_space = {'x_min': -500, 'y_max': 210, 'y_min': -17}
    
    def image_recv(self):
        ret, frame = self.cap.read()
        return frame

    def get_pose(self, print_flag=False):
        time.sleep(0.1)
        ret = self.robot.get_tcp_position()
        if ret[0] == 0:
            if print_flag:
                print("the eef position is :", ret[1])
            return ret[1]
        else:
            print("some things happend,the errcode is: ", ret[0])
            return [0] * 6

    def move(self, target_pose, speed=90, acc=30, tol=2):
        global ERROR
        # ret_end_pose = self.robot.get_tcp_position()[1]
        # ret = self.robot.linear_move(end_pos=ret_end_pose, move_mode=ABS, is_block=True, speed=speed)
        # ret = self.robot.linear_move(end_pos=target_pose, move_mode=ABS, is_block=True, speed=speed)

        # while True:
        #     cur_pos = np.array(self.get_pose())
        #     tar_pose = np.array(target_pose)

        #     pose_error = abs(cur_pos - tar_pose)
        #     if (pose_error < 0.5).all():
        #         break
        #     else:
        #         ret = self.robot.linear_move(end_pos=target_pose, move_mode=ABS, is_block=True, speed=speed)  #
        #         print('===retry move to target pose', ERROR)
        #         ERROR += 1
        #         continue

        source_pose = self.get_pose(print_flag=True)

        if (abs(np.array(source_pose) - np.array(target_pose)) < 0.5).all():
            return

        for _ in range(10):
            
            ret = self.robot.linear_move_extend(end_pos=target_pose, move_mode=ABS, is_block=True, speed=speed, acc=acc, tol=tol)
            time.sleep(0.5)
            if ret == (0,):
                return 
            else:
                print('move error code: ', ret[0])
            
            
        
        if ret[0] != 0:
            print('lastest error code', self.robot.get_last_error())
            print('clear error code')
            self.robot.clear_error()
            print('recover source pose')
            ret = self.robot.linear_move_extend(end_pos=source_pose, move_mode=ABS, is_block=True, speed=speed, acc=acc, tol=tol)
            return



    def catch_and_place(self, test=False):
        self.move(pose_home)
        # _ = self.get_pose(print_flag=True)
        self.gripper.open()
        self.move(pose_cup)
        _ = self.get_pose(print_flag=True)
        self.gripper.pose(55)  # 55
        self.move(pose_cup_pre)
        # _ = self.get_pose(print_flag=True)
        self.move(pose_coffee_pre)
        self.move(pose_coffee)
        
        time.sleep(5)

    def go_home(self):
        self.move(pose_home)

    def move_to_car(self, place_choice_id, test=False):
        global pose_first_place
        global detector_tag
        if(place_choice_id == 1):
            
            self.move(pose_coffee_pre)
            self.move(pose_place_pre)    
            time.sleep(8)
            self.detect_holder(place_choice_id-1, detector_tag)
            

            print('-------------------')
            time.sleep(3)
            self.move(pose_detect_init)
            self.pose_refine(detector_tag, max_move_y=200, max_move_x=200)
            print('pose refine finished')

            cur_pos = self.get_pose()
            pose_first_place = deepcopy(cur_pos)

            
            cur_pos[2] = 50
            self.move(cur_pos)
            time.sleep(1)
            now_pos = self.get_pose()

            self.gripper.open()

            time.sleep(2)
            # return
            self.move(pose_home)
            time.sleep(2)
        elif(place_choice_id == 2):
            self.move(pose_coffee_pre)
            self.move(pose_place_pre)
            time.sleep(3)
            print('-------------------')
            self.move(pose_detect_init)
            pose_first_place[1] = min(pose_first_place[1]+140,220)
            self.move(pose_first_place)

            pose_first_place[2] = 50
            self.move(pose_first_place)
            time.sleep(1)
            now_pos = self.get_pose()

            self.gripper.open()

            time.sleep(2)
            # return
            self.move(pose_home)
            time.sleep(2)

        # del detector_tag

    def put_to_plate(self):
        self.move(pose_coffee_pre)
        self.move(pose_place_pre) 
        self.move(put_pre)
        self.move(put)
        time.sleep(1)
        self.gripper.open()
        self.move(put_pre)
        self.move(pose_place_pre)
        self.move(pose_home)
        return True

    def detect_holder(self, choice, detector_tag):
        #place_choice = [[-0.08, 0.07], [0.08, 0.09], [0, 0.07]]
        place_choice = [[-0.05, 0.02], [0.05, 0.02]]
        print('放在托盘{}号位置_{}'.format(choice+1, place_choice[choice]))
        #choice = np.random.randint(0, 2)
        #choice = place_choice_id
        num = 200
        while num:
            num -= 1
            print('111111111111111')
            frame = deepcopy(self.image_recv())
            # ret, frame = self.cap.read()
            print('22222222222222')
            while len(frame) == 0:
                frame = deepcopy(self.image_recv())
                print('Try to reload image')
            # frame = deepcopy(self.image_recv())

            print('333333333333333333')
            detector_tag.tag_size = 0.028
            detected_tag = detector_tag.detect_apriltag(frame)
            print('444444444444444444')
            for tag in detected_tag:
                if tag.tag_id == 7:
                    print('==============tag 7 detected=============')
                    self.holder_pos = [tag.pose_t[0], tag.pose_t[1]]
                    # tag_pose2place_pose
                    print('holder_pos change pre:',self.holder_pos)
                    self.holder_pos[0] += place_choice[choice][0]
                    self.holder_pos[1] += place_choice[choice][1]
                    print('holder_pos change latter:',self.holder_pos)
                else:
                    print('=============tag 7 NOT detected=============')
                    self.holder_pos = -1
            print('55555555555555555555555')
            if self.holder_pos != -1:
                break
            else:
                print('Try to detect holder_tag, again. ')
                cv2.imwrite('holder.png', frame)
        if self.holder_pos == -1:
            print('Tag detect ERROR!!')
            exit()
        print('holder_pose', self.holder_pos)

    def _try_move(self, pose):
        print('Try to move randomly...')
        random_x = (np.random.rand() - 0.5) * 50
        random_y = (np.random.rand() - 0.5) * 50
        random_z = (np.random.rand() - 0.5) * 40
        pose[0] += random_x
        pose[1] += random_y
        pose[2] += random_z
        return pose

    def pose_refine(self, detector_tag, max_move_x=10, max_move_y=30, P=0.526, debug=False):
        try:
            if isinstance(self.holder_pos, list):
                pass
        except:
            print('not list')
            exit()

        while True:
            unusual_flag = False
            detected_cup = None
            first = True
            detector_tag.tag_size = 0.02
            # 如果检测不到tag，让机械臂随机移动，尝试改变视角检测到tag
            ori_pose = self.get_pose()
            while not detected_cup:
                frame = deepcopy(self.image_recv())
                while len(frame) == 0:
                    frame = deepcopy(self.image_recv())
                    print('Try to reload image')

                detected_tag = detector_tag.detect_apriltag(frame)
                if len(detected_tag) == 0:
                    print('no tag')
                    continue
                for tag in detected_tag:
                    if tag.tag_id == 10:
                        detected_cup = [tag.pose_t[0], tag.pose_t[1]]
                        detected_cup[1] -= 0.1  # tag to cup
                # 避免超出工作区间，一直以第一次的位置为中心，尝试移动

                if not first:
                    try_pose = self._try_move(deepcopy(ori_pose))
                    NUM[1] += 1
                    self.move(try_pose, speed=30)
                    time.sleep(1)
                first = False

            cur_pose = self.get_pose()
            next_pose = deepcopy(cur_pose)

            dc = deepcopy(detected_cup)
            print("cup", dc)
            # dc[0] -= 320
            # dc[0] = dc[0] * 36 /40.0
            # dc[1] = dc[1] * 36 /40.0
            # max_move_x=10, max_move_y=30
            move_x = min(max_move_x, (self.holder_pos[1] - dc[1]) * 1000)
            move_y = min(max_move_y, (self.holder_pos[0] - dc[0]) * 1000)

            # move_x = self.holder_pos[1] - dc[1] * 1000
            # move_y = self.holder_pos[0] - dc[0] * 1000

            # move_x = move_x.item()
            # move_y = move_y.item()
            next_pose[0] = next_pose[0] + move_x/5
            next_pose[1] = next_pose[1] + move_y/6
            # next_pose[1] = next_pose[1] + move_y
            print('-------------------------------')
            print('摄像头检测到的托盘位置',self.holder_pos)
            print('摄像头检测到的杯子位置:', dc)

            print('计算出来的需要移动的move_x:',move_x)
            print('计算出来的需要移动move_y:',move_y)
            
            print('机械臂读取到的当前位置: ', cur_pose)
            print('计算出来的机械臂要移动的位置:', next_pose)
            
            if next_pose[1] > self.work_space['y_max']:
                next_pose[1] = self.work_space['y_max']
                NUM[0] += 1
                # if DEBUG:
                print('trigger y_max')
                unusual_flag = True
            if next_pose[1] < self.work_space['y_min']:
                next_pose[1] = self.work_space['y_min']
                NUM[0] += 1
                print('trigger y_min')
                unusual_flag = True            
            if next_pose[0] < self.work_space['x_min']:
                next_pose[0] = self.work_space['x_min']
                NUM[0] += 1
                print('trigger x_min')
                unusual_flag = True
            if DEBUG:
                print("move_x: ", move_x, "move_y", move_y)
                print("next_pose", next_pose[:3])


            


            print('移动')
            self.move(next_pose, speed=50)

            if abs(move_x) < 15 and abs(move_y) < 15 and not unusual_flag:
                print("Detect place_ pose, successfully!\n")
                break
            time.sleep(2)

            # break
        # self.holder_pos = -1

    def close(self):
        try:
            self.gripper.close()
            self.robot.logout()
        except BaseException as e:
            pass

    # def test_detect_tag(self):
    #     frame = deepcopy(self.image_recv())
    #     detected_tag = self.detector_tag.detect_apriltag(frame)
    #     time.sleep(1)

# pos of grasping cup
pose_cup = [-349.34854, 122.30499, 389.59625, 1.5069654215532, 0.7686112244748222, -1.8903693967144113]
# initial pos of robot arm
pose_home = deepcopy(pose_cup)
pose_home[0] += 150
# the previous pos of grasping cup
pose_cup_pre = deepcopy(pose_cup)
pose_cup_pre[0] += 120
# pos of coffee machine

pose_coffee = [-372.14708, -207.15975, 303.73204995137866, 1.5677426719928, 0.7954727138694333, -0.956521397053689]
pose_coffee_pre = deepcopy(pose_coffee)
pose_coffee_pre[0] *= 0.75
pose_coffee_pre[1] *= 0.75
# place coffee on the blue robot
pose_place_pre = [-217.7743219562775, 120.80791028140658, 315, 1.5681723065969577, 0.7828422200612521,
                  -1.9134782644303079]
# pose_place_pre = [-217.7743219562775, 77.80791028140658, 315, 1.5681723065969577, 0.7828422200612521,
#                   -1.9134782644303079]
#-217
pose_detect_init = [-444.0724984508521, 100.93983999288257, 180.78272150569, 1.6027293600799686, 0.845561659417912,
                    -1.5703964629902765]
# pose_detect_init = [-444.0724984508521, 56.93983999288257, 180.78272150569, 1.6027293600799686, 0.845561659417912,
#                     -1.5703964629902765]
#56 180
pose_place = [-502.2573407638289, 202.50320623481147, 43.62471457067907, 1.5925724689142353, 0.7521966879969761,
              -1.5562127312492604]
pose_first_place = [0,0,0,0,0,0]
# put cup to plate
put_pre = [-341.5986791322935, 104.9173080098673, 148.0674933383293, 1.7626193797153014, 0.8378390425932971, -1.397733578702065]
put = [-341.5986791322934, 104.91730800986734, 89.50880071137482, 1.7626193797153005, 0.8378390425932953, -1.397733578702064] 
if __name__ == '__main__':
    # Gripper.init()
    jaka = Jaka(jaka_ip='192.168.1.133')
    # jaka.holder_pos = [385.0, 201.0]
    # jaka.move(pose_detect_init)
    # for i in range(5):
    #     jaka.get_pose(print_flag=True)
    #     time.sleep(1)
    # while 1:
    #     jaka.test_detect_tag()
    # DEBUG = True
    # # jaka.get_pose(print_flag=True)
    # # jaka.move(pose_detect_init)
    # jaka.detect_holder()
    # jaka.get_pose(print_flag=True)
    # for i in range(100):
    #     print('The ', i, ' start:')
    #     jaka.move(pose_detect_init)


    # # jaka.holder_pos= [0.07056428, -0.00126146]
    #     jaka.pose_refine(max_move_y=200, max_move_x=200)
    #     cur_pos = jaka.get_pose()
    #     cur_pos[2] = 43
    #     jaka.move(cur_pos)
    #     time.sleep(3)
    # # time.sleep(2)
    # # jaka.move(pose_home)
    # print('fail times:', NUM[0])
    # print('Try times:', NUM[1])
    # # print(pose_place_pre)
    # jaka.close()
